from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
import torchvision.transforms as T
import torch
import requests
from io import BytesIO

def predict_image(url_or_path, device="cpu"):
    # 1. 加载 processor & model
    processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
    model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224").to(device)
    model.eval()

    # 2. 载入并预处理图像
    if url_or_path.startswith("http"):
        img = Image.open(BytesIO(requests.get(url_or_path).content)).convert("RGB")
    else:
        img = Image.open(url_or_path).convert("RGB")

    transform = T.Compose([
        T.Resize((224, 224)),
        T.ToTensor(),
        T.Normalize(mean=processor.image_mean, std=processor.image_std),
    ])
    pixel_values = transform(img).unsqueeze(0).to(device)  # [1,3,224,224]
    
    # 3. 端到端前向
    with torch.no_grad():
        outputs = model(pixel_values=pixel_values)
        logits = outputs.logits  # [1,1000]
        probs = torch.softmax(logits, dim=-1)
        top5 = torch.topk(probs, 5, dim=-1)

    # 4. 打印 Top-5
    for idx, score in zip(top5.indices[0], top5.values[0]):
        print(f"{model.config.id2label[idx.item()]}: {score.item():.4f}")
    print("Final:", model.config.id2label[top5.indices[0][0].item()])

# 用法示例
predict_image("https://raw.githubusercontent.com/microsoft/ComputerVision/master/images/cat.jpg",
              device="cuda:0" if torch.cuda.is_available() else "cpu")
